#include "co.h"
#include <stdlib.h>
#include <string.h>
#include <stdlib.h>
#include <setjmp.h>
#include <stdint.h>
#include <stdio.h>
#include <assert.h>

#define STACK_SIZE (64 * 1024)

enum co_status {
  CO_NEW = 1, // 新创建，还未执行过
  CO_RUNNING, // 已经执行过
  CO_WAITING, // 在 co_wait 上等待
  CO_DEAD,    // 已经结束，但还未释放资源
};

typedef struct co {
  char name[32];
  void (*func)(void *); // co_start 指定的入口地址和参数
  void *arg;

  struct co     *prev, *next;
  struct co     *waiter;  // 是否有其他协程在等待当前协程
  enum co_status status;  // 协程的状态
  jmp_buf        context; // 寄存器现场 (setjmp.h)
  uint8_t        stack[STACK_SIZE] __attribute__ ((aligned (16))); // 协程的堆栈
} Coroutine;

static Coroutine *co_current = NULL;
static Coroutine *co_main = NULL;

static void stack_switch_call(void *sp, void *entry, uintptr_t arg) {
  asm volatile (
#if __x86_64__
    "movq %0, %%rsp; movq %2, %%rdi; jmp *%1"
      : : "b"((uintptr_t)sp - 8),     "d"(entry), "a"(arg)
#else
    "movl %0, %%esp; movl %2, 4(%0); jmp *%1"
      : : "b"((uintptr_t)sp - 16), "d"(entry), "a"(arg)
#endif
  );
}

__attribute__((constructor)) static void co_init() {
	co_main = (Coroutine *)malloc(sizeof(Coroutine));
	assert(co_main);
	memset(co_main, 0, sizeof(Coroutine));
	strcpy(co_main->name, "main");
	co_main->next = co_main->prev = co_main;
	co_main->status = CO_RUNNING;
	co_current = co_main;
}

static void co_free(Coroutine *co) {
	assert(co);
	Coroutine *prev = co->prev;
	Coroutine *next = co->next;
	prev->next = next;
	next->prev = prev;

	free(co);
}

static void co_handler(Coroutine *co) {
	assert(co);
	(co->func)(co->arg);
	co->status = CO_DEAD;
	if (co->waiter) {
		co->waiter->status = CO_RUNNING;
		co->waiter = NULL;
	}
	co_yield();
}

static Coroutine *co_schedule() {
	assert(co_main);
	Coroutine *co = co_current->next;
	while (co->status == CO_DEAD || co->status == CO_WAITING) { co = co->next; }
	return co;
}

struct co *co_start(const char *name, void (*func)(void *), void *arg) {
	assert(co_main);
	Coroutine *co = (Coroutine *)malloc(sizeof(Coroutine));
	assert(co);
	memset(co, 0, sizeof(Coroutine));
	strcpy(co->name, name);
	co->func = func;
	co->arg = arg;
	co->status = CO_NEW;

	Coroutine *tail = co_main->prev;
	tail->next = co;
	co_main->prev = co;
	co->next = co_main;
	co->prev = tail;

  return co;
}

/* 表示当前协程需要等待，直到 co 协程的执行完成才能继续执行 */
void co_wait(struct co *co) {
	assert(co && co_current);
	if (co->status == CO_DEAD) { co_free(co); return; }
	co_current->status = CO_WAITING;
	co->waiter = co_current;
	co_yield();
}

void co_yield() {
	int val = setjmp(co_current->context);
	if (val == 0) {
		co_current = co_schedule();
		if (co_current->status == CO_NEW) {
			co_current->status = CO_RUNNING;
			stack_switch_call(&co_current->stack[STACK_SIZE], co_handler, (uintptr_t)co_current);
		} else {
			longjmp(co_current->context, 0);
		}
	}
}
